package tree.medium;

import java.util.LinkedList;
import java.util.Queue;

/**
 * 键值映射
 * 实现一个 MapSum 类，支持两个方法，insert 和 sum：
 *
 * MapSum() 初始化 MapSum 对象
 * void insert(String key, int val) 插入 key-val 键值对，字符串表示键 key ，整数表示值 val 。如果键 key 已经存在，那么原来的键值对将被替代成新的键值对。
 * int sum(string prefix) 返回所有以该前缀 prefix 开头的键 key 的值的总和。
 *
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode-cn.com/problems/map-sum-pairs
 */
public class MapSumPairs_677 {

    class TrieNode{
        int count;
        TrieNode[] tns = new TrieNode[26];
    }
    TrieNode root;

    public MapSumPairs_677(){
        root = new TrieNode();
    }
    public void insert(String key, int val) {
        TrieNode cur = root;

        for (int i = 0; i < key.length(); i++) {
            int index = key.charAt(i) - 'a';
            if (cur.tns[index] == null) {
                cur.tns[index] = new TrieNode();
            }
            cur = cur.tns[index];
        }
        cur.count = val;
    }


    public int sum(String prefix) {
        TrieNode cur = root;
        for (int i = 0; i < prefix.length(); i++) {
            int index = prefix.charAt(i) - 'a';
            if (cur.tns[index] == null) {
                return 0;
            }
            cur = cur.tns[index];
        }
        return dfs(cur);
    }

    public int dfs(TrieNode node){
        int ans = 0;
        Queue<TrieNode> queue = new LinkedList<>();
        queue.offer(node);
        while (!queue.isEmpty()) {
            TrieNode poll = queue.poll();
            ans += poll.count;
            for (int i = 0; i < poll.tns.length; i++) {
                if (poll.tns[i] == null){
                    continue;
                }
                queue.offer(poll.tns[i]);
            }
        }
        return ans;
    }

    public static void main(String[] args) {

    }
}
